# Copyright 2020-2021 The Kubeflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from kfp import components
from kfp import dsl

component_op_1 = components.load_component_from_text(
    """
name: upstream
inputs:
- {name: input_1, type: String}
- {name: input_2, type: Float}
- {name: input_3, type: }
- {name: input_4}
- {name: input_5, type: Metrics}
- {name: input_6, type: Datasets}
- {name: input_7, type: Some arbitrary type}
- {name: input_8, type: {GcsPath: {data_type: TSV}}}
outputs:
- {name: output_1, type: Integer}
- {name: output_2, type: Model}
- {name: output_3}
implementation:
  container:
    image: gcr.io/google.com/cloudsdktool/cloud-sdk:latest
    command:
    - bash
    - -exc
    args:
    - |
        input_1=$0
        input_2=$1
        input_3_uri=$2
        input_4_uri=$3
        input_5_uri=$4
        input_6_uri=$5
        input_7_uri=$6
        input_8_uri=$7
        output_1_path=$8
        output_2_uri=$9
        output_3_path=${10}

        echo "input_1: ${input_1}"
        echo "input_2: ${input_2}"
        echo "input_3: $(gsutil cat "${input_3_uri}")"
        echo "input_4: $(gsutil cat "${input_4_uri}")"
        echo "input_5: $(gsutil cat "${input_5_uri}")"
        echo "input_6: $(gsutil cat "${input_6_uri}")"
        echo "input_7: $(gsutil cat "${input_7_uri}")"
        echo "input_8: $(gsutil cat "${input_8_uri}")"
        mkdir -p "$(dirname "${output_1_path}")"
        echo -n "1" > "${output_1_path}"
        echo "MODEL_PAYLOAD" | gsutil cp - "${output_2_path}"
        mkdir -p "$(dirname "${output_3_path}")"
        echo -n "output_3_value" > "${output_3_path}"


    - {inputValue: input_1}
    - {inputValue: input_2}
    - {inputUri: input_3}
    - {inputUri: input_4}
    - {inputUri: input_5}
    - {inputUri: input_6}
    - {inputUri: input_7}
    - {inputUri: input_8}
    - {outputPath: output_1}
    - {outputUri: output_2}
    - {outputPath: output_3}
"""
)

component_op_2 = components.load_component_from_text(
    """
name: downstream
inputs:
- {name: input_a, type: Integer}
- {name: input_b, type: Model}
- {name: input_c}
implementation:
  container:
    image: gcr.io/google.com/cloudsdktool/cloud-sdk:latest
    command:
    - bash
    - -exc
    args:
    - |
        input_a="$0"
        input_b_uri="$1"
        input_c_path="$2"
        echo "input_a: ${input_a}"
        echo "input_b: $(gsutil cat "${input_b_uri}")"
        echo "input_c: $(cat "${input_c_path}")"
    - {inputValue: input_a}
    - {inputUri: input_b}
    - {inputPath: input_c}
"""
)


@dsl.pipeline(name='pipeline-with-various-types')
def my_pipeline(
    input1: str,
    input3,
    input4='',
    input5='gs://bucket/metrics',
    input6='gs://bucket/dataset',
    input7='arbitrary value',
    input8='gs://path2'
):
    component_1 = component_op_1(
        input_1=input1,
        input_2=3.1415926,
        input_3=input3,
        input_4=input4,
        input_5='gs://bucket/metrics',
        input_6=input6,
        input_7=input7,
        input_8=input8
    )
    component_2 = component_op_2(
        input_a=component_1.outputs['output_1'],
        input_b=component_1.outputs['output_2'],
        input_c=component_1.outputs['output_3']
    )
